from fastapi import File, UploadFile, HTTPException, Body
from fastapi.responses import StreamingResponse, JSONResponse

from core.config import router

from service.qa_service import process_uploaded_file, answer_question
import os
import json
from core.config import vectorstore


@router.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    if not os.getenv("DASHSCOPE_API_KEY"):
        raise HTTPException(status_code=400, detail="请设置DASHSCOPE_API_KEY环境变量")

    global vectorstore
    file_content = await file.read()
    # 传入文件名用于多文档区分
    new_vectorstore, chunks = process_uploaded_file(file_content, file.filename)

    if not new_vectorstore:
        raise HTTPException(status_code=400, detail="文件处理出错")

    # 合并向量存储（增量添加）
    if vectorstore is None:
        vectorstore = new_vectorstore
    else:
        vectorstore.add_documents(chunks)

    return {
        "message": "文件处理成功",
        "filename": file.filename,
        "chunk_count": len(chunks)
    }


@router.get("/document_chunks")
async def get_document_chunks(target_filename: str = None):
    if not vectorstore:
        raise HTTPException(status_code=400, detail="向量存储未初始化，请先上传文件")

    # 从向量数据库中获取文档分块
    if target_filename:
        docs = vectorstore.get(where={"filename": target_filename})
    else:
        docs = vectorstore.get()  # 获取所有文档

    # 在get_document_chunks接口中，对结果去重
    chunks = []
    seen = set()  # 用于记录已处理的文档内容+文件名，避免重复
    for i, doc in enumerate(docs['documents']):
        key = (doc, docs['metadatas'][i].get("filename"))
        if key not in seen:
            seen.add(key)
            chunks.append({
                "id": i,
                "filename": docs['metadatas'][i].get("filename"),
                "content": doc,
                "metadata": docs['metadatas'][i]
            })
    return JSONResponse({"chunks": chunks})


@router.post("/question_and_answer")
async def question_and_answer(
        question: str = Body(..., embed=True),
        target_filename: str = Body(None, embed=True)  # 可选：指定文档提问
):
    """流式问答接口，支持指定文档检索"""
    if not os.getenv("DASHSCOPE_API_KEY"):
        raise HTTPException(status_code=400, detail="请设置DASHSCOPE_API_KEY环境变量")

    if not vectorstore:
        raise HTTPException(status_code=400, detail="向量存储未初始化，请先上传文件")

    # 转发生成器内容为SSE事件流
    def event_stream():
        for item in answer_question(
                vectorstore,
                question,
                target_filename  # 传入目标文档名
        ):
            print(item['data'])
            yield f"event: {item['type']}\r\ndata: {json.dumps(item['data'], ensure_ascii=False)}\r\n\r\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")
